/**
 * Software License, Version 1.0
 * 
 * Copyright 2003 The Trustees of Indiana University.  All rights reserved.
 * 
 *
 *Redistribution and use in source and binary forms, with or without 
 *modification, are permitted provided that the following conditions are met:
 *
 *1) All redistributions of source code must retain the above copyright notice,
 * the list of authors in the original source code, this list of conditions and
 * the disclaimer listed in this license;
 *2) All redistributions in binary form must reproduce the above copyright 
 * notice, this list of conditions and the disclaimer listed in this license in
 * the documentation and/or other materials provided with the distribution;
 *3) Any documentation included with all redistributions must include the 
 * following acknowledgement:
 *
 *"This product includes software developed by the Community Grids Lab. For 
 * further information contact the Community Grids Lab at 
 * http://communitygrids.iu.edu/."
 *
 * Alternatively, this acknowledgement may appear in the software itself, and 
 * wherever such third-party acknowledgments normally appear.
 * 
 *4) The name Indiana University or Community Grids Lab or NaradaBrokering, 
 * shall not be used to endorse or promote products derived from this software 
 * without prior written permission from Indiana University.  For written 
 * permission, please contact the Advanced Research and Technology Institute 
 * ("ARTI") at 351 West 10th Street, Indianapolis, Indiana 46202.
 *5) Products derived from this software may not be called NaradaBrokering, 
 * nor may Indiana University or Community Grids Lab or NaradaBrokering appear
 * in their name, without prior written permission of ARTI.
 * 
 *
 * Indiana University provides no reassurances that the source code provided 
 * does not infringe the patent or any other intellectual property rights of 
 * any other entity.  Indiana University disclaims any liability to any 
 * recipient for claims brought by any other entity based on infringement of 
 * intellectual property rights or otherwise.  
 *
 *LICENSEE UNDERSTANDS THAT SOFTWARE IS PROVIDED "AS IS" FOR WHICH NO 
 *WARRANTIES AS TO CAPABILITIES OR ACCURACY ARE MADE. INDIANA UNIVERSITY GIVES
 *NO WARRANTIES AND MAKES NO REPRESENTATION THAT SOFTWARE IS FREE OF 
 *INFRINGEMENT OF THIRD PARTY PATENT, COPYRIGHT, OR OTHER PROPRIETARY RIGHTS. 
 *INDIANA UNIVERSITY MAKES NO WARRANTIES THAT SOFTWARE IS FREE FROM "BUGS", 
 *"VIRUSES", "TROJAN HORSES", "TRAP DOORS", "WORMS", OR OTHER HARMFUL CODE.  
 *LICENSEE ASSUMES THE ENTIRE RISK AS TO THE PERFORMANCE OF SOFTWARE AND/OR 
 *ASSOCIATED MATERIALS, AND TO THE PERFORMANCE AND VALIDITY OF INFORMATION 
 *GENERATED USING SOFTWARE.
 */
package cgl.narada.benchmark.security;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Calendar;
import java.util.GregorianCalendar;
import java.util.Hashtable;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;

import org.apache.log4j.Logger;

import cgl.narada.event.EventHeaders;
import cgl.narada.event.NBEvent;
import cgl.narada.event.NBEventException;
import cgl.narada.event.TemplateProfileAndSynopsisTypes;
import cgl.narada.service.ServiceException;
import cgl.narada.service.client.ClientService;
import cgl.narada.service.client.EventProducer;
import cgl.narada.service.client.SessionService;
import cgl.narada.service.qos.impl.QosServiceImpl;
import cgl.narada.service.security.SecureTopicsManager;
import cgl.narada.service.security.impl.EntityOperationsImpl;
import cgl.narada.service.security.kmc.SecurityToken;
import cgl.narada.service.security.kmc.SignedSecurityToken;
import cgl.narada.service.security.kmc.TSSGenerator;
import cgl.narada.service.security.kmc.TimeStampedSeqNo;
import cgl.narada.service.security.kmc.TopicRights;
import cgl.narada.service.security.securityprovider.CertificateManager;
import cgl.narada.service.security.securityprovider.CertificateUtil;
import cgl.narada.service.timer.ClockI;
import cgl.narada.service.timer.HRClock;
import cgl.narada.util.SystemInit;
import cgl.narada.util.SystemNotInitializedException;

/**
 * <br>
 * Created on Apr 10, 2006
 * 
 * @author Harshawardhan Gadgil (hgadgil@grids.ucs.indiana.edu)
 */
public class OperationsTest {

    public static void main(String[] args) {

        SystemInit.init();

        int entityId = 1;

        Hashtable entityIDTimeStampMapping = new Hashtable();
        TimeStampedSeqNo veryOldTimeStamp = TSSGenerator.generate();
        entityIDTimeStampMapping.put(new Integer(entityId), veryOldTimeStamp);

        String configFilePath = null;
        try {
            configFilePath = SystemInit
                    .getInitData(SystemInit.SYS_SERVICE_CONFIG_PATH);
        } catch (SystemNotInitializedException e3) {
            // TODO Auto-generated catch block
            e3.printStackTrace();
        }

        CertificateManager certMan = new CertificateManager();
        String keystore = null;
        try {
            keystore = SystemInit.getInitData(SystemInit.SYS_KEYSTORE);
        } catch (SystemNotInitializedException e3) {
            // TODO Auto-generated catch block
            e3.printStackTrace();
        }
        certMan.init(keystore, "passpass");

        Certificate clientCert = CertificateUtil.getCertificate(certMan, "c1");
        PrivateKey privKey = CertificateUtil.getPrivateKey(certMan, "c1");
        PublicKey pubKey = CertificateUtil.getPublicKey(certMan, "c1");
        PublicKey rootPubKey = CertificateUtil.getPublicKey(certMan, "rootca");

        Certificate kmcCert = CertificateUtil.getCertificate(certMan, "kmc1");
        PrivateKey kmcPrivKey = CertificateUtil.getPrivateKey(certMan, "kmc1");

        String nbHome = null;
        try {
            nbHome = SystemInit.getInitData(SystemInit.SYS_NB_HOME);
        } catch (SystemNotInitializedException e3) {
            // TODO Auto-generated catch block
            e3.printStackTrace();
        }

        ClockI hrclock = HRClock.getClock(nbHome);

        KeyGenerator kg = null;
        try {
            kg = KeyGenerator.getInstance("AES");
        } catch (NoSuchAlgorithmException e2) {
            // TODO Auto-generated catch block
            e2.printStackTrace();
        }

        kg.init(256, new SecureRandom());
        SecretKey secretKey = kg.generateKey();

        // Set the default properties...
        String mode = "CBC", padding = "PKCS7Padding", signingAlgorithm = "SHA1withRSA";

        int size = 16 * 1024;

        byte[] originalPayload = new byte[size];

        // Signed security token
        Calendar validUntil = GregorianCalendar.getInstance();
        TopicRights topicRights = new TopicRights(TopicRights.PUBLISH_RIGHT);

        SecurityToken token = new SecurityToken(clientCert, topicRights,
                validUntil);

        SignedSecurityToken sst = new SignedSecurityToken(token, kmcCert);
        sst.generateSignedDigest(kmcPrivKey);

        byte[] sstBytes = sst.getBytes();

        // Create a producer
        EventProducer producer = null;
        EntityOperationsImpl eoi = null;
        try {
            ClientService clientService = SessionService.getClientService(
                    entityId, configFilePath);

            producer = clientService.createEventProducer();
            eoi = new EntityOperationsImpl();

        } catch (ServiceException e1) {
            // TODO Auto-generated catch block
            e1.printStackTrace();
        }

        FileOutputStream fout = null;
        try {
            fout = new FileOutputStream("/home/hgadgil/OperationsTest-1.out");
        } catch (FileNotFoundException e1) {
            // TODO Auto-generated catch block
            e1.printStackTrace();
        }

        // ////////////////////////////////////////////////////
        // ------------------ Publish phase Benchmark ------------
        // /////////////////////////////////////////////////////

        long start = 0, t1 = 0, t2 = 0, t3 = 0, t4 = 0, t5 = 0, t6 = 0;
        NBEvent nbEvent = null;

        for (int i = 0; i < 1; i++) {
            System.out.println("Publisher ---> " + i);
            try {

                // TEST 1: Publishing event
                byte[] encryptedPayload = null;
                byte[] iv = null;
                start = hrclock.getTimeMicroseconds();
                // Step 1: Encrypt payload

                String algModePadding = secretKey.getAlgorithm() + "/" + mode
                        + "/" + padding;

                for (int run = 0; run < 100; run++) {
                    iv = eoi.getIV(secretKey.getAlgorithm(), mode, padding);
                }

                // BM:
                t1 = (hrclock.getTimeMicroseconds() - start) / 100;

                start = hrclock.getTimeMicroseconds();

                for (int run = 0; run < 100; run++) {
                    encryptedPayload = eoi.encryptPayload(secretKey,
                            originalPayload, algModePadding, iv, "BC");
                }
                // BM:
                t2 = (hrclock.getTimeMicroseconds() - start) / 100;

                // Step 2: Sign encrypted payload

                byte[] signature = null;

                try {
                    EntityOperationsImpl entityOperations = new EntityOperationsImpl();

                    // BM:
                    start = hrclock.getTimeMicroseconds();

                    for (int run = 0; run < 100; run++) {

                        signature = entityOperations.signPayload(privKey,
                                signingAlgorithm, "BC", encryptedPayload);
                    }
                    // BM:
                    t3 = (hrclock.getTimeMicroseconds() - start) / 100;

                } catch (ServiceException e) {}

                // replace the content payload
                nbEvent = producer.generateEvent(
                        TemplateProfileAndSynopsisTypes.STRING, "topicName",
                        encryptedPayload);

                // nbEvent.setContentPayload(encryptedPayload);

                // Set the security token and signature.
                EventHeaders headers = nbEvent.getEventHeaders();

                headers.setEncryptedPayloadInformation(iv, mode, padding);
                headers.setSecurityToken(sstBytes);
                headers.setSignatureInformation(signingAlgorithm, signature);

                TimeStampedSeqNo stssn = TSSGenerator.generate();

                if (stssn == null) {
                    System.err
                            .println("Unable to get timestamp, Cannot route packet !");
                    return;
                }

                headers.setSecurityTimeStampedSequenceNo(stssn.getBytes());

                // Step 3: Sign event header
                // ---------------------
                // Sign the Event header and include that
                byte[] headerSignature = null;

                try {
                    EntityOperationsImpl entityOperations = new EntityOperationsImpl();

                    // BM:
                    start = hrclock.getTimeMicroseconds();

                    for (int run = 0; run < 100; run++) {
                        headerSignature = entityOperations.signPayload(privKey,
                                signingAlgorithm, "BC", headers.getBytes());
                    }

                    // BM:
                    t4 = (hrclock.getTimeMicroseconds() - start) / 100;

                    nbEvent.getEventProperties().setImmutableProperty(
                            SecureTopicsManager.HEADER_SIGNATURE,
                            headerSignature, QosServiceImpl.class.getName());
                } catch (NBEventException e1) {
                    // TODO Auto-generated catch block
                    e1.printStackTrace();
                }

                try {
                    fout
                            .write(("" + t1 + "\t" + t2 + "\t" + t3 + "\t" + t4 + "\n")
                                    .getBytes());
                } catch (IOException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                }

            } catch (ServiceException e) {
                // TODO Auto-generated catch block
                e.printStackTrace();
            }

        }

        try {
            fout.close();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        // /////--------------------------------------------

        // ////////////////////////////////////////////////////
        // ------------------ Broker Verification phase Benchmark ------------
        // /////////////////////////////////////////////////////

        try {
            fout = new FileOutputStream("/home/hgadgil/OperationsTest-3.out");
        } catch (FileNotFoundException e1) {
            // TODO Auto-generated catch block
            e1.printStackTrace();
        }

        for (int i = 0; i < 105; i++) {
            System.out.println("Broker ---> " + i);

            // Verify Event
            start = hrclock.getTimeMicroseconds();

            for (int run = 0; run < 100; run++) {
                boolean validated = verifyNBEvent(nbEvent, rootPubKey);
            }

            t1 = (hrclock.getTimeMicroseconds() - start) / 100;

            // Check for replay by comparing timestamps
            start = hrclock.getTimeMicroseconds();

            for (int run = 0; run < 100; run++) {
                // Check the time stamp for this sender
                Integer eId = new Integer(nbEvent.getEventHeaders().getSource());
                byte[] stssnBytes = nbEvent.getEventHeaders()
                        .getSecurityTimeStampedSeqNo();
                TimeStampedSeqNo stssn = TimeStampedSeqNo.unmarshal(stssnBytes);

                if (stssn == null) {
                    // Something is wrong with this event...Discard it...
                    log.error("Secure Event does not contain "
                            + "Time Stamped Sequence No. !");
                }

                // System.out.println("TimeStamp: " + stssn.toString());

                TimeStampedSeqNo lastTimeStamp = (TimeStampedSeqNo) entityIDTimeStampMapping
                        .get(eId);

                if (lastTimeStamp == null) {
                    // This is the first event from this guy.. so simply store
                    // his timestamp...
                    // System.out.println("New Timestamp Entry !");

                    entityIDTimeStampMapping.put(eId, stssn);
                } else {
                    // Check if the timestamp in the event is monotonically
                    // increasing...
                    if (stssn.isGreater(lastTimeStamp)) {
                        // Good..replace lasttimestamp with the new timestamp
                        entityIDTimeStampMapping.put(new Integer(2), stssn);
                    } else {
                        // Aaha.. Duplicate time stamp... discard event...
                        log.error("STSSN<" + nbEvent.getContentSynopsis()
                                + "> is " + "<= last known timestamp");
                    }
                }
            }

            t2 = (hrclock.getTimeMicroseconds() - start) / 100;

            SignedSecurityToken sst2 = SignedSecurityToken
                    .createObjectFromBytes(nbEvent.getEventHeaders()
                            .getSignedSecurityToken());
            start = hrclock.getTimeMicroseconds();
            for (int run = 0; run < 100; run++) {
                sst2.getSecurityToken().isValid();
            }

            t3 = (hrclock.getTimeMicroseconds() - start) / 100;

            try {
                fout.write(("" + t1 + "\t" + t2 + "\t" + t3 + "\n").getBytes());
            } catch (IOException e) {
                e.printStackTrace();
            }

        }

        try {
            fout.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        System.exit(0);

        // /////--------------------------------------------

        // ////////////////////////////////////////////////////
        // ------------------ Subscribe phase Benchmark ------------
        // /////////////////////////////////////////////////////

        // TEST 2: Validating and routing event

        try {
            fout = new FileOutputStream("/home/hgadgil/OperationsTest-2.out");
        } catch (FileNotFoundException e1) {
            e1.printStackTrace();
        }

        EventHeaders header = nbEvent.getEventHeaders();

        // Verify the Signed security token's certificate
        // SignedSecurityToken sst = SignedSecurityToken
        // .createObjectFromBytes(header.getSignedSecurityToken());

        EntityOperationsImpl entityOperations = null;
        try {
            entityOperations = new EntityOperationsImpl();
        } catch (ServiceException e2) {
            // TODO Auto-generated catch block
            e2.printStackTrace();
        }
        Certificate entityCert = sst.getSecurityToken()
                .getTopicOwnerCertificate();

        for (int i = 0; i < 105; i++) {

            System.out.println("Subscriber: --> " + i);

            try {

                // BM:
                start = hrclock.getTimeMicroseconds();

                for (int run = 0; run < 100; run++) {

                    // Step 1: Verify payload (SST, certificate etc...)

                    // Step 1-1: Verify the certificate with Root CA's public
                    // key
                    Certificate cert = sst.getKmcCertificate();
                    if (!CertificateUtil.ValidateCertificate(cert, rootPubKey)) {
                        System.err
                                .println("KMC's certificate could not be validated !");
                        return;
                    }

                    // Step 1-2: Now check the payload's signature
                    // Step 1-2a: Get the entity's certificate from the
                    // certificate's table

                    // Verify the certificate
                    if (!CertificateUtil.ValidateCertificate(entityCert,
                            rootPubKey)) {
                        System.err
                                .println("Entity's certificate could not be validated !");
                        return;
                    }

                    // Validate the Signed Security token
                    if (!sst.verify()) {
                        System.err
                                .println("Could not verify integrity of signed security token !");
                        return;
                    }

                    // Check Header signature
                    byte[] sig;
                    try {
                        sig = (byte[]) nbEvent.getEventProperties()
                                .getProperty(
                                        SecureTopicsManager.HEADER_SIGNATURE);
                    } catch (NBEventException e) {
                        sig = null;
                    }

                    if (sig == null) {
                        System.err.println("Header signature not found !");
                        return;
                    }

                    // Check the signature

                    if (!entityOperations.validateSignature(entityCert
                            .getPublicKey(), signingAlgorithm, "BC", header
                            .getBytes(), sig)) {
                        System.err
                                .println("Could not validate Event header signature !");
                        return;
                    }
                }
                t1 = (hrclock.getTimeMicroseconds() - start) / 100;

                // Check payload signature
                // Step 2: Check Signature

                // BM:
                start = hrclock.getTimeMicroseconds();

                for (int run = 0; run < 100; run++) {
                    if (!entityOperations.validateSignature(entityCert
                            .getPublicKey(), signingAlgorithm, "BC", nbEvent
                            .getContentPayload(), header.getSignature())) {
                        System.err
                                .println("Could not validate Event payload signature !");
                        return;
                    }
                }

                // BM:
                t2 = (hrclock.getTimeMicroseconds() - start) / 100;

                // Ok, Everything checks FINE... this event is good to
                // go !

                // Step 3: Decrypt

                originalPayload = nbEvent.getContentPayload();

                byte[] iv = header.getCipherIV();

                String algModePadding = secretKey.getAlgorithm() + "/" + mode
                        + "/" + padding;

                start = hrclock.getTimeMicroseconds();
                byte[] decryptedPayload;
                for (int run = 0; run < 100; run++) {

                    decryptedPayload = eoi.decryptPayload(secretKey,
                            originalPayload, algModePadding, iv, "BC");
                }
                // BM:
                t3 = (hrclock.getTimeMicroseconds() - start) / 100;

                try {
                    fout.write(("" + t1 + "\t" + t2 + "\t" + t3 + "\n")
                            .getBytes());
                } catch (IOException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                }

            } catch (ServiceException e1) {
                // TODO Auto-generated catch block
                e1.printStackTrace();
            }
        }

        try {
            fout.close();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    private static Hashtable revokedCertificates = new Hashtable();
    static Logger log = Logger.getLogger("OperationsTest");

    private static boolean verifyNBEvent(NBEvent event,
            PublicKey rootCAPublicKey) {

        EventHeaders headers = event.getEventHeaders();

        log.info("Verifying NBEvent !!");

        if (!headers.hasSignature()) {
            log.error("No security token found !");
            return false;
        }

        // Verify the Signed security token's certificate
        SignedSecurityToken sst = SignedSecurityToken
                .createObjectFromBytes(headers.getSignedSecurityToken());

        // Step 1: Verify the certificate with Root CA's public key
        Certificate cert = sst.getKmcCertificate();
        if (!CertificateUtil.ValidateCertificate(cert, rootCAPublicKey)) {
            log.error("KMC's certificate could not be validated !");
            return false;
        }

        // Step 2: Now check the payload's signature
        // Step 2a: Get the entity's certificate from the certificate's table
        Certificate entityCert = sst.getSecurityToken()
                .getTopicOwnerCertificate();

        // TODO: Is there a better way to deal with revoked Certificates ?
        BigInteger certId = ((X509Certificate) entityCert).getSerialNumber();

        // Check if the certificate is in the revoked list
        if (revokedCertificates.contains(certId)) {
            log.error("Entity's certificate has been revoked !");
            return false;
        }

        // Verify the certificate
        if (!CertificateUtil.ValidateCertificate(entityCert, rootCAPublicKey)) {
            log.error("Entity's certificate could not be validated !");
            return false;
        }

        // Validate the Signed Security token
        if (!sst.verify()) {
            log.error("Could not verify integrity of signed security token !");
            return false;
        }

        // Check the signature
        try {
            EntityOperationsImpl entityOperations = new EntityOperationsImpl();

            // Check Header signature
            byte[] sig;
            try {
                sig = (byte[]) event.getEventProperties().getProperty(
                        SecureTopicsManager.HEADER_SIGNATURE);
            } catch (NBEventException e) {
                sig = null;
            }

            if (sig == null) {
                log.error("Header signature not found !");
                return false;
            }

            String signingAlgorithm = headers.getSignatureAlgorithm();

            if (!entityOperations.validateSignature(entityCert.getPublicKey(),
                    signingAlgorithm, "BC", headers.getBytes(), sig)) {
                log.error("Could not validate Event header signature !");
                return false;
            }

            // Check payload signature
            if (!entityOperations.validateSignature(entityCert.getPublicKey(),
                    signingAlgorithm, "BC", event.getContentPayload(), headers
                            .getSignature())) {
                log.error("Could not validate Event payload signature !");
                return false;
            }
        } catch (ServiceException e) {
            log.error("", e);
            return false;
        }

        // Finally check the signature of the headers and see if they were
        // modifies

        // Ok, Everything checks FINE... this event is good to go !
        return true;
    }
}
